package sam.learn.framework.dynamicDataSource;

import cn.hutool.core.util.StrUtil;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.annotation.DS;
import com.baomidou.dynamic.datasource.support.DataSourceClassResolver;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import sam.learn.biz.user.repo.entity.UserInfo;
import sam.learn.biz.user.service.UserInfoService;
import sam.learn.framework.security.CurrentUser;

import javax.sql.DataSource;
import java.util.Objects;

public class DynamicDataSourceAnnotationInterceptor implements MethodInterceptor {

    private final DataSourceClassResolver dataSourceClassResolver;

    private final CurrentUser currentUser;
    private final UserInfoService userInfoService;
    private final DynamicRoutingDataSource dynamicRoutingDataSource;
    private final DbDynamicDataSourceLoader dbDynamicDataSourceLoader;

    public DynamicDataSourceAnnotationInterceptor(
            CurrentUser currentUser, UserInfoService userInfoService
            , DynamicRoutingDataSource dynamicRoutingDataSource
            , DbDynamicDataSourceLoader dbDynamicDataSourceLoader
    ) {
        dataSourceClassResolver = new DataSourceClassResolver(true);
        this.currentUser = currentUser;
        this.userInfoService = userInfoService;
        this.dynamicRoutingDataSource = dynamicRoutingDataSource;
        this.dbDynamicDataSourceLoader = dbDynamicDataSourceLoader;
    }

    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        String dsKey = determineDatasourceKey(invocation);
        checkDatasource(dsKey);
        DynamicDataSourceContextHolder.push(dsKey);
        try {
            return invocation.proceed();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    /**
     * Determine the key of the datasource
     *
     * @param invocation MethodInvocation
     * @return dsKey
     */
    private String determineDatasourceKey(MethodInvocation invocation) {
        String key = dataSourceClassResolver.findKey(invocation.getMethod(), invocation.getThis(), DS.class);
        if (StrUtil.isBlank(key)) {
            String username = currentUser.currentUser().getUsername();
            UserInfo userInfo = userInfoService.queryUserInfoBy(username);
            return userInfo.getRegionCode();
        }
        return key;
    }

    private void checkDatasource(String regionCode) {
        DataSource dataSource = dynamicRoutingDataSource.getDataSource(regionCode);
        if (Objects.nonNull(dataSource)) {
            dbDynamicDataSourceLoader.loadDatasourceBy(regionCode);
        }
    }
}
